#ifndef _WIN32_WINNT		// Allow use of features specific to Windows XP or later.                   
#define _WIN32_WINNT 0x0501	// Change this to the appropriate value to target other versions of Windows.
#endif	

#include <windows.h>
#include "NTHook.h"
#include "hwbrk.h"



HANDLE CNTHook::mHVH = NULL;
DWORD CNTHook::mKiSystemCall = 0;
DWORD CNTHook::mKiSystemCallRet = 0;

CNTHook::CNTHook(void)
{
    mAryFun[ZwGetContextThread] = 0x55;
    mAryFun[ZwSetContextThread] = 0xD5;
    mAryFun[ZwQueryVirtualMemory] = 0xB2;
    mAryFun[ZwReadVirtualMemory] = 0xBA;
    mAryFun[ZwWriteVirtualMemory] = 0x115;
    mHcall= mHret = NULL;
}

CNTHook::~CNTHook(void)
{
    if (mHVH)
    {
        RemoveVectoredExceptionHandler(mHVH);
        mHVH = 0;
    }
}

bool CNTHook::LoadHandler()
{
    mHVH = AddVectoredExceptionHandler(1, CNTHook::VectoredHandler);
    return mHVH!=0;
}

bool CNTHook::Hook(HANDLE hThread)
{
    DWORD KiSystemCall, KiSystemCallRet;
    __asm
    {
        push    eax;
        mov     eax, 0x7FFE0300;
        mov     eax, dword ptr[eax];
        mov     KiSystemCall, eax;
        add     eax, 4;
        mov     KiSystemCallRet, eax;
        pop     eax;
    }

    {
    //8B D4 0F 34 C3
        BYTE* pTest = (BYTE*) KiSystemCall;
        if (*pTest != 0x8B) return false;
        mKiSystemCall = KiSystemCall, mKiSystemCallRet = KiSystemCallRet;
    }

    mHcall = SetHardwareBreakpoint(hThread, HWBRK_TYPE_CODE, HWBRK_SIZE_1, (void*)KiSystemCall);
    mHret = SetHardwareBreakpoint(hThread, HWBRK_TYPE_CODE, HWBRK_SIZE_1, (void*)KiSystemCallRet);
    return mHcall && mHret;
}

bool CNTHook::UnHook(HANDLE hThread)
{
    bool bret = 0;
    bret = RemoveHardwareBreakpoint(mHret);
    bret &= RemoveHardwareBreakpoint(mHcall);
    return bret;
}

LONG WINAPI CNTHook::VectoredHandler(PEXCEPTION_POINTERS ExceptionInfo)
{
#define PC ExceptionInfo->ContextRecord

    if (PC->Eip == mKiSystemCall)
    {
        PC->Edx = PC->Esp;
        PC->Eip ++;
        MessageBox(0,L"in",0,0);
        return EXCEPTION_CONTINUE_EXECUTION;;
    }
    else if (PC->Eip == mKiSystemCallRet)
    {
        PC->Edx = PC->Esp;
        PC->Eip ++;
        MessageBox(0,L"ret",0,0);
        return EXCEPTION_CONTINUE_EXECUTION;;
    }

    MessageBox(0,L"sb",0,0);
    return EXCEPTION_CONTINUE_SEARCH;
#undef PC
}